//
// Created by Alex on 17/3/13.
//


#include "socket.h"

static void SetAddr(struct sockaddr_in &addr, const char *ip, const unsigned short port) {
  bzero(&addr, sizeof(addr));
  addr.sin_family = AF_INET;
  addr.sin_port = htons(port);

  if (ip == nullptr
      || '\0' == *ip
      || 0 == strcmp(ip, "0")
      || 0 == strcmp(ip, "0.0.0.0")
      || 0 == strcmp(ip, "*")) {
    addr.sin_addr.s_addr = htonl(INADDR_ANY);

  } else {
    addr.sin_addr.s_addr = inet_addr(ip);
  }
}

static int SetNonBlocking(int fd, bool enable) {
  int flags = fcntl(fd, F_GETFL, 0);
  if (enable) {
    flags |= O_NONBLOCK;
  } else {
    flags &= ~O_NONBLOCK;
  }
  return fcntl(fd, F_SETFL, flags);
}

static bool parseAddr(const std::string &conn,
                      std::string &type, std::string &ip, int &port) {
  size_t pos = 0;
  if ((pos = conn.find_first_of("tcp://")) != std::string::npos) {
    type = "tcp";
    pos += strlen("tcp://");
  } else if ((pos = conn.find_first_of("udp://")) != std::string::npos) {
    type = "udp";
    pos += strlen("udp://");
  } else {
    return false;
  }

  size_t split = conn.find_first_of(':', pos);
  if (split != std::string::npos) {
    ip = conn.substr(pos, split - pos);
    port = std::atoi(conn.substr(split + 1).c_str());
  }
  return !type.empty() && !ip.empty() && port > 0;
}

Socket::Socket(const std::string &conn) :
    m_socketFd(-1),
    m_recvTimeout(0),
    m_sendTimeout(0),
    m_connectTimeout(0),
    m_isLinger(true),
    m_lingerValue(1),
    m_noDelay(true) {
  parseAddr(conn, m_protocol, m_peerIp, m_peerPort);
}

Socket::Socket(int fd) :
    m_socketFd(fd),
    m_peerPort(0),
    m_sendTimeout(0),
    m_connectTimeout(0),
    m_recvTimeout(0),
    m_isLinger(true),
    m_lingerValue(1),
    m_noDelay(true),
    m_refCnt(0) {

  if (m_socketFd < 0) {
    return;
  }
  SetNonBlocking(m_socketFd, true);

  sockaddr_in addr;
  socklen_t len = sizeof(addr);
  int ret = getsockname(m_socketFd, (struct sockaddr *) &addr, &len);
  if (ret != -1) {
    m_localIp = inet_ntoa(addr.sin_addr);
    m_localPort = ntohs(addr.sin_port);
  }
}


Socket::~Socket() {
  Close();
}

bool Socket::IsOpen() {
  return m_socketFd != -1;
}

void Socket::SetLinger(bool on, int linger) {
  m_isLinger = on;
  m_lingerValue = linger;
}

void Socket::SetInterrupt(std::shared_ptr<bool> interrupt) {
  m_interrupt = interrupt;
}

bool Socket::Peek() {
  if (!IsOpen()) {
    return false;
  }

  if (m_interrupt != nullptr && *m_interrupt == false) {
    return false;
  }

  int event = kWriteEvent | kReadEvent;
  bool pollRet = cr_poll(m_socketFd, event, m_recvTimeout);
  if (!pollRet) {
    return false;
  }

  uint8_t buf;
  ssize_t r = recv(m_socketFd, (char *) &buf, 1, MSG_PEEK);
  if (r == -1) {
    Close();
  }
  return r > 0;
}

bool Socket::OpenConnection() {
  if (IsOpen()) {
    return true;
  }

  m_socketFd = -1;
  if (m_protocol == "tcp") {
    m_socketFd = socket(AF_INET, SOCK_STREAM, 0);
    if (m_socketFd != -1) {
      SetLinger(m_isLinger, m_lingerValue);
      SetNoDelay(m_noDelay);
    }

  } else if (m_protocol == "udp") {
    m_socketFd = socket(AF_INET, SOCK_DGRAM, 0);
  }

  if (m_socketFd == -1) {
    return false;
  }

  SetNonBlocking(m_socketFd, true);

  sockaddr_in addr;
  SetAddr(addr, m_peerIp.c_str(), m_peerPort);
  int ret = cr_connect(m_socketFd, (struct sockaddr *) &addr, sizeof(addr), m_connectTimeout);
  if (ret != 0) {
    return false;
  }

  socklen_t len = sizeof(addr);
  memset(&addr, 0, sizeof(addr));
  ret = getsockname(m_socketFd, (struct sockaddr *) &addr, &len);
  if (ret != -1) {
    m_localIp = inet_ntoa(addr.sin_addr);
    m_localIp = ntohs(addr.sin_port);
  }
  return true;
}

bool Socket::Open() {
  if (IsOpen()) {
    return true;
  }
  return OpenConnection();
}

void Socket::Close() {
  if (m_socketFd != -1) {
    shutdown(m_socketFd, SHUT_RDWR);
    close(m_socketFd);
    m_socketFd = -1;
  }
}

void Socket::SetNoDelay(bool noDelay) {
  int flag = noDelay ? 1 : 0;
  setsockopt(m_socketFd, IPPROTO_TCP, TCP_NODELAY, (char *) &flag, sizeof(flag));
}


void Socket::SetRecvTimeout(int ms) {
  m_recvTimeout = ms;
}

void Socket::SetSendTimeout(int ms) {
  m_sendTimeout = ms;
}

uint32_t Socket::Write(const uint8_t *buf, uint32_t len) {
  if (!IsOpen()) {
    return 0;
  }
  if (m_interrupt != nullptr && *m_interrupt) {
    return 0;
  }

  ssize_t ret = cr_write(m_socketFd, buf, len, m_sendTimeout);
  if (ret > 0) {
    return ret;
  }
  Close();
  return 0;
}

uint32_t Socket::Read(uint8_t *buf, uint32_t len) {
  if (!IsOpen()) {
    return 0;
  }
  if (m_interrupt != nullptr && *m_interrupt) {
    return 0;
  }

  ssize_t ret = cr_read(m_socketFd, buf, len, m_recvTimeout);
  if (ret > 0) {
    return ret;
  }
  Close();
  return 0;
}

void Socket::SetIp(const std::string &ip) {
  m_peerIp = ip;
}

void Socket::SetPort(int port) {
  m_peerPort = port;
}

void Socket::Ref() {
  m_refCnt += 1;
}

bool Socket::UnRef() {
  m_refCnt -= 1;
  if (m_refCnt <= 0) {
    delete this;
    return true;
  }
  return false;
}

ServerSocket::ServerSocket(const std::string &ip, int port) :
    m_listening(false),
    m_ip(ip),
    m_port(port),
    m_socketFd(-1),
    m_acceptBacklog(DEFAULT_BACKLOG),
    m_sendTimeout(0),
    m_recvTimeout(0),
    m_tcpRecvBuffer(0),
    m_tcpSendBuffer(0),
    m_listenCallback(nullptr),
    m_acceptCallback(nullptr),
    m_interrupt(new bool) {
  *m_interrupt = false;
}

ServerSocket::~ServerSocket() {
  Close();
}

void ServerSocket::SetSendTimeout(int sendTimeout) {
  m_sendTimeout = sendTimeout;
}

void ServerSocket::SetRecvTimeout(int recvTimeout) {
  m_recvTimeout = recvTimeout;
}

void ServerSocket::SetAccpetCallback(const socket_accept_fun_t &callback) {
  m_acceptCallback = callback;
}

void ServerSocket::SetListenCallback(const socket_listen_fun_t &callback) {
  m_listenCallback = callback;
}

void ServerSocket::SetAcceptBackLog(int val) {
  m_acceptBacklog = val;
}

void ServerSocket::SetTcpRecvBuffer(int kb) {
  m_tcpRecvBuffer = kb * 1024;
}

void ServerSocket::SetTcpSendBuffer(int kb) {
  m_tcpSendBuffer = kb * 1024;
}

void ServerSocket::Close() {
  if (m_socketFd > 0) {
    shutdown(m_socketFd, SHUT_RDWR);
    close(m_socketFd);
    m_socketFd = -1;
  }
  m_listening = false;

  if (m_interrupt != nullptr) {
    *m_interrupt = true;
  }
}

bool ServerSocket::Listen() {
  m_socketFd = -1;
  if (m_port < 0 || m_port > 0xFFFF) {
    return false;
  }
  m_socketFd = socket(AF_INET, SOCK_STREAM, 0);
  m_listening = true;
  *m_interrupt = false;

  int one = 1;
  //send buffer size
  if (m_tcpSendBuffer > 0) {
    setsockopt(m_socketFd, SOL_SOCKET, SO_SNDBUF,
               (char *) &m_tcpSendBuffer, sizeof(m_tcpSendBuffer));
  }

  //recv buffer size
  if (m_tcpRecvBuffer > 0) {
    setsockopt(m_socketFd, SOL_SOCKET, SO_RCVBUF,
               (char *) &m_tcpRecvBuffer, sizeof(m_tcpRecvBuffer));
  }

  //defer accept
  setsockopt(m_socketFd, IPPROTO_TCP, TCP_DEFER_ACCEPT, (char *) &one, sizeof(one));

  //linger
  struct linger ling = {0, 0};
  setsockopt(m_socketFd, SOL_SOCKET, SO_LINGER, (char *) &ling, sizeof(ling));

  //no delay
  setsockopt(m_socketFd, IPPROTO_TCP, TCP_NODELAY, (char *) &one, sizeof(one));

  //non blocking
  SetNonBlocking(m_socketFd, true);

  //reuse addrr and port
  int flags = 1;
  setsockopt(m_socketFd, SOL_SOCKET, SO_REUSEADDR, &flags, sizeof(flags));

  struct sockaddr_in addr;
  SetAddr(addr, m_ip.c_str(), m_port);

  int ret = bind(m_socketFd, (struct sockaddr *) &addr, sizeof(addr));
  if (ret != 0) {
    close(m_socketFd);
    return false;
  }

  if (m_listenCallback != nullptr) {
    m_listenCallback(m_socketFd);
  }

  ret = listen(m_socketFd, m_acceptBacklog);
  if (ret == -1) {
    Close();
  }
  return ret == 0;
}

void ServerSocket::Interrupt() {
  if (m_interrupt != nullptr) {
    *m_interrupt = true;
  }
}

Socket *ServerSocket::Accept() {
  if (m_socketFd < 0 || !m_listening || *m_interrupt) {
    return nullptr;
  }
  sockaddr_in cliAddr;
  socklen_t size = sizeof(sockaddr_in);
  int clientFd = cr_accept(m_socketFd, (struct sockaddr *) &cliAddr, &size);

  if (clientFd > 0) {
    //set non blocking
    SetNonBlocking(clientFd, true);

    Socket *socket = new Socket(clientFd);
    //set children interrupt
    socket->SetInterrupt(m_interrupt);

    //send timeout
    socket->SetRecvTimeout(m_recvTimeout);
    socket->SetSendTimeout(m_sendTimeout);

    //set socket info
    socket->SetIp(inet_ntoa(cliAddr.sin_addr));
    socket->SetPort(ntohs(cliAddr.sin_port));

    //callback
    if (m_acceptCallback != nullptr) {
      m_acceptCallback(socket);
    }
    return socket;
  }
  return nullptr;
}


